#!/usr/bin/env python

import rospy
from mavros_msgs.msg import PositionTarget, State
from mavros_msgs.srv import SetMode, CommandBool
from std_msgs.msg import Header

# 全局变量，用于存储无人机状态
current_state = State()

# 回调函数，更新无人机状态
def state_callback(state):
    global current_state
    current_state = state

# 初始化ROS节点
rospy.init_node('drone_takeoff')

# 创建发布者，发布到/mavros/setpoint_raw/local话题
setpoint_pub = rospy.Publisher('/mavros/setpoint_raw/local', PositionTarget, queue_size=10)

# 创建订阅者，订阅/mavros/state话题以获取无人机状态
state_sub = rospy.Subscriber('/mavros/state', State, state_callback)

# 等待MAVROS服务启动
rospy.loginfo("等待MAVROS服务启动...")
rospy.wait_for_service('/mavros/cmd/arming')
rospy.wait_for_service('/mavros/set_mode')
rospy.loginfo("MAVROS服务已启动")

# 创建服务代理，用于设置模式和解锁
set_mode_client = rospy.ServiceProxy('/mavros/set_mode', SetMode)
arming_client = rospy.ServiceProxy('/mavros/cmd/arming', CommandBool)

# 设置目标高度
target_altitude = 100.0  # 目标高度为3米

# 创建PositionTarget消息
target_pose = PositionTarget()

# 设置目标位置的相关参数
target_pose.header = Header()
target_pose.header.frame_id = "base_footprint"  # 设置坐标系的标识符
target_pose.coordinate_frame = PositionTarget.FRAME_LOCAL_NED  # 使用局部NED坐标系

# 设置掩码，忽略速度、加速度和偏航角，仅控制位置
target_pose.type_mask = (
    PositionTarget.IGNORE_VX |  # 忽略速度X
    PositionTarget.IGNORE_VY |  # 忽略速度Y
    PositionTarget.IGNORE_VZ |  # 忽略速度Z
    PositionTarget.IGNORE_AFX |  # 忽略加速度X
    PositionTarget.IGNORE_AFY |  # 忽略加速度Y
    PositionTarget.IGNORE_AFZ |  # 忽略加速度Z
    PositionTarget.IGNORE_YAW |  # 忽略偏航角
    PositionTarget.IGNORE_YAW_RATE  # 忽略偏航角速率
)

# 设置目标位置 (x, y, z)
target_pose.position.x = 0.0
target_pose.position.y = 0.0
target_pose.position.z = target_altitude  # 目标高度

# 设置发布频率
rate = rospy.Rate(10)  # 10Hz

# 等待与飞控建立连接
rospy.loginfo("等待与飞控建立连接...")
while not rospy.is_shutdown() and not current_state.connected:
    rate.sleep()
rospy.loginfo("已与飞控建立连接")

# 在解锁前，持续发布目标位置
rospy.loginfo("在解锁前，持续发布目标位置...")
for i in range(50):  # 发布50次，确保飞控接收足够的目标位置
    target_pose.header.stamp = rospy.Time.now()
    setpoint_pub.publish(target_pose)
    rate.sleep()

# 切换到OFFBOARD模式
rospy.loginfo("尝试切换到OFFBOARD模式...")
offboard_set_mode = SetMode()
offboard_set_mode.custom_mode = "OFFBOARD"
for i in range(10):  # 尝试10次
    if set_mode_client(custom_mode="OFFBOARD").mode_sent:
        rospy.loginfo("已切换到OFFBOARD模式")
        break
    rate.sleep()
else:
    rospy.logerr("无法切换到OFFBOARD模式")
    exit(1)

# 解锁无人机
rospy.loginfo("尝试解锁无人机...")
for i in range(10):  # 尝试10次
    if arming_client(value=True).success:
        rospy.loginfo("无人机已解锁")
        break
    rate.sleep()
else:
    rospy.logerr("无法解锁无人机")
    exit(1)

# 发布目标位置，控制无人机起飞
rospy.loginfo("发布目标位置，控制无人机起飞...")
while not rospy.is_shutdown():
    target_pose.header.stamp = rospy.Time.now()  # 更新时间戳
    setpoint_pub.publish(target_pose)
    rate.sleep()
